import albumentations as A
from torch.utils.data import DataLoader

import numpy as np

from medfmc.datasets.medical_datasets import MultiLabelDataset

class Calibrator(MultiLabelDataset):
    def __init__(self, dataset, batch_size=4, **kwargs):
        super().__init__(**kwargs)
        self.dataset = dataset
        self.batch_size = batch_size


    def distribution_calibration(self, query, base_means, base_cov, k=2, alpha=0.21):
        dist = []
        for i in range(len(base_means)):
            dist.append(np.linalg.norm(query-base_means[i]))
        index = np.argpartition(dist, k)[:k]
        mean = np.concatenate([np.array(base_means)[index], query[np.newaxis, :]])
        calibrated_mean = np.mean(mean, axis=0)
        calibrated_cov = np.mean(np.array(base_cov)[index], axis=0)+alpha

        return calibrated_mean, calibrated_cov

    def resample_augment_data(self):
        # Define the augmentation pipeline
        augmentation = A.Compose([
            # Add your desired augmentation techniques here
            A.HorizontalFlip(p=0.5),
            A.RandomBrightnessContrast(p=0.2),
            # Add more augmentation techniques as needed
        ])

        # Create a new list to store augmented samples
        augmented_data = []

        # Resample and augment each sample in the dataset
        for data in self.dataset:
            image = data['img_info']['filename']
            label = data['gt_label']

            # Apply augmentation to the image
            augmented_image = augmentation(image=image)['image']

            # Create a new data dictionary for the augmented sample
            augmented_info = {
                'img_prefix': data['img_prefix'],
                'img_info': {'filename': augmented_image},
                'gt_label': label
            }

            augmented_data.append(augmented_info)

        # Create a DataLoader for the augmented dataset
        augmented_loader = DataLoader(augmented_data, batch_size=self.batch_size, shuffle=True)

        return augmented_loader

# Usage example
if __name__ == '__main__':
    # Load the Chest19 dataset
    chest19_dataset = Chest19(ann_file='path/to/annotations.txt', data_prefix='path/to/images/')

    # Create an instance of DataAugmentation
    calibrator = CalibrateDistribution(chest19_dataset)

    # Resample and augment the dataset
    augmented_loader = calibrator.resample_augment_data()

    # Iterate over the augmented data loader
    for batch in augmented_loader:
        # Perform further processing or training using the augmented data
        # ...
        print(batch)